import os
import sys
import argparse
import numpy as np
import mindspore

from mindspore import Tensor,export,load_checkpoint,context
from mind3d.models.dgcnn import DGCNN_seg,DGCNN_cls
from mind3d.utils.load_yaml import load_yaml
from mind3d.models.pointnet2 import Pointnet2segModelMSG,Pointnet2segModelSSG


if not os.path.exists("./mindir"):
    os.mkdir("./mindir")

num_classes=40
part_num=50

def dgcnn_export(opt,args):
    classifier=DGCNN_seg(opt, opt['train'].get("k"))
    load_checkpoint(args.model,net=classifier)
    input_data=Tensor(np.random.uniform(0.0,1.0,size=[3,4096,9]),mindspore.float32)
    export(classifier,input_data,file_name="./mindir/dgcnn_seg_model_6",file_format=args.file_format)
    print("successfully export model")

def pointnet2_export(opt,args):
    num_segm = 20
    segmentation = Pointnet2segModelSSG(num_classes=num_segm, use_color=False, use_normal=False)
    load_checkpoint(args.model, net=segmentation)
    input_data = Tensor(np.random.uniform(0.0, 1.0, size=[4, 1024, 3]), mindspore.float32)
    input_label = Tensor(np.random.uniform(0.0, 1.0, size=[4, 1024, 1]), mindspore.float32)
    export(segmentation, input_data, file_name='./mindir/pointnet2_seg', file_format=args.file_format)
    print("successfully export model")





if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Export.")
    parser.add_argument("--device_target", default="GPU", help='device')
    parser.add_argument('--model', type=str, default="/home/czy/HuaWei/ms3d/example/pointnet2/pretrained_ckpt/350.ckpt", help="model path")
    parser.add_argument("--file_format", type=str, default="MINDIR", help="export file format")
    parser.add_argument("-opt", type=str, default="/home/czy/HuaWei/final2023/mind3d/configs/pointnet2/pointnet2_scannet_seg.yaml",
                        help='Path to option YAML file.')
    args = parser.parse_known_args()[0]
    os.environ["CUDA_VISBLE_DEVICES"] = '1'
    context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target,  max_call_depth=20480)

    opt = load_yaml(args.opt)
    network_name = opt["model_name"]
    if network_name == "DGCNN_seg":
        dgcnn_export(opt, args)
    elif network_name == "Pointnet2_seg":
        pointnet2_export(opt, args)



